import logging
import os

from llama_index.core import VectorStoreIndex, Settings
from llama_index.core.schema import TextNode
from llama_index.core.vector_stores import MetadataFilters, MetadataFilter  # 新增导入
from llama_index.embeddings.huggingface import HuggingFaceEmbedding
from llama_index.llms.dashscope import DashScope

logging.basicConfig(level=logging.DEBUG)
print('初始化模型对象')
llm_dash = DashScope(
    model="qwen-turbo",  # 使用通义千问Max模型 qwen-max
    api_key=os.getenv("DASHSCOPE_API_KEY"),  # 从阿里云控制台获取
    temperature=0.3
)
embed_model_zh = HuggingFaceEmbedding(
    model_name="BAAI/bge-small-zh-v1.5",  # 中文优化模型
    device="cpu"
)
# 配置全局设置
Settings.embed_model = embed_model_zh
Settings.llm = llm_dash

nodes = [
    TextNode(
        text="一起事件是指意外或恶意的事故，它有可能对我们的IT资产安全造成不良影响。",
        metadata={"department": "Security"},
    ),
    TextNode(
        text="一起事件是IT服务的意外中断或性能下降。这些都会导致IT资产缩水",
        metadata={"department": "IT"},
    )
]

# 修正1：直接构建索引并启用元数据存储
vector_index = VectorStoreIndex(
    nodes=nodes,
    store_metadata=True  # 启用元数据存储
)

# 修正2：创建元数据过滤器
filters_param = MetadataFilters(
    filters=[
        MetadataFilter(key="department", value="Security")  # 过滤IT部门
    ]
)

# 应用过滤器到检索器
retriever = vector_index.as_retriever(filters=filters_param)  # 传入过滤器

result = retriever.retrieve("IT资产")
print('检索结果:')
for item in result:
    print(f"==> {item.text} [元数据: {item.metadata}]")
